/*
 * Tencent is pleased to support the open source community by making BK-CI 蓝鲸持续集成平台 available.
 *
 * Copyright (C) 2019 Tencent.  All rights reserved.
 *
 * BK-CI 蓝鲸持续集成平台 is licensed under the MIT license.
 *
 * A copy of the MIT License is included in this file.
 *
 *
 * Terms of the MIT License:
 * ---------------------------------------------------
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of
 * the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
 * LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
 * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package com.tencent.bkrepo.common.mongo.reactive.dao

import com.tencent.bkrepo.common.mongo.api.util.MongoDaoHelper
import com.tencent.bkrepo.common.mongo.api.util.MongoDaoHelper.determineShardingCount
import com.tencent.bkrepo.common.mongo.api.util.MongoDaoHelper.determineShardingFields
import com.tencent.bkrepo.common.mongo.api.util.MongoDaoHelper.shardingValues
import com.tencent.bkrepo.common.mongo.api.util.MongoDaoHelper.shardingValuesOf
import com.tencent.bkrepo.common.mongo.api.util.sharding.ShardingUtils
import com.tencent.bkrepo.common.mongo.util.MongoIndexResolver
import jakarta.annotation.PostConstruct
import kotlinx.coroutines.reactor.awaitSingle
import org.slf4j.LoggerFactory
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.beans.factory.annotation.Value
import org.springframework.data.domain.Page
import org.springframework.data.domain.PageImpl
import org.springframework.data.domain.PageRequest
import org.springframework.data.mongodb.core.ReactiveMongoTemplate
import org.springframework.data.mongodb.core.aggregation.Aggregation
import org.springframework.data.mongodb.core.index.IndexDefinition
import org.springframework.data.mongodb.core.query.Query
import java.lang.reflect.Field

abstract class ShardingMongoReactiveDao<E> : AbstractMongoReactiveDao<E>() {

    @Suppress("LateinitUsage")
    @Autowired
    lateinit var reactiveMongoTemplate: ReactiveMongoTemplate

    @Value("\${sharding.count:#{null}}")
    private val fixedShardingCount: Int? = null

    /**
     * 分表Field，key为列名
     */
    protected lateinit var shardingFields: LinkedHashMap<String, Field>

    /**
     * 分表数
     */
    protected var shardingCount: Int = 1

    /**
     * 分表工具类
     */
    protected val shardingUtils by lazy {
        determineShardingUtils()
    }

    @PostConstruct
    private fun init() {
        this.shardingFields = determineShardingFields(classType, customShardingColumns())
        this.shardingCount = determineShardingCount(classType, shardingUtils, customShardingCount())
        ensureIndex()
    }

    private fun ensureIndex() {
        if (shardingCount < 0) {
            return
        }
        val start = System.currentTimeMillis()
        val indexDefinitions = MongoIndexResolver.resolveIndexFor(classType)
        val nonexistentIndexDefinitions = filterExistedIndex(indexDefinitions)
        nonexistentIndexDefinitions.forEach {
            for (i in 1..shardingCount) {
                val reactiveMongoOperations = determineReactiveMongoOperations()
                val collectionName = parseSequenceToCollectionName(i - 1)
                reactiveMongoOperations.indexOps(collectionName).ensureIndex(it)
                    .subscribe { indexName ->
                        logger.info("Ensure index [$indexName] for sharding collection.")
                    }
            }
        }

        val indexCount = shardingCount * indexDefinitions.size
        val consume = System.currentTimeMillis() - start

        logger.info("Ensure [$indexCount] index for sharding collection [$collectionName], consume [$consume] ms.")
    }

    protected open fun customShardingColumns(): List<String> {
        return emptyList()
    }

    protected open fun customShardingCount(): Int? {
        return fixedShardingCount
    }

    override fun determineReactiveMongoOperations(): ReactiveMongoTemplate {
        return reactiveMongoTemplate
    }

    override fun determineCollectionName(query: Query): String {
        val shardingValues = shardingValuesOf(query.queryObject, shardingFields)
        requireNotNull(shardingValues) { "Sharding value can not empty!" }
        return shardingKeyToCollectionName(shardingValues)
    }

    override fun determineCollectionName(entity: E): String {
        return shardingKeyToCollectionName(shardingValues(entity as Any, shardingFields))
    }

    override fun determineCollectionName(aggregation: Aggregation): String {
        val shardingValues = shardingValuesOf(aggregation, shardingFields)
        require(!shardingValues.isNullOrEmpty()) { "Sharding values can not be empty!" }
        return shardingKeyToCollectionName(shardingValues)
    }

    override fun determineCollectionName(): String {
        return MongoDaoHelper.determineShardingCollectionName(classType) ?: super.determineCollectionName()
    }

    /**
     * 支持查询条件不包含sharding key的分页查询
     *
     * @param pageRequest 分页信息
     * @param query 查询条件，不要在其中包含分页查询条件
     */
    suspend fun pageWithoutShardingKey(pageRequest: PageRequest, query: Query): Page<E> {
        if (shardingCount <= 0 || shardingCount > MAX_SHARDING_COUNT_OF_PAGE_QUERY) {
            throw UnsupportedOperationException()
        }

        val startIndex = pageRequest.pageNumber * pageRequest.pageSize
        var limit = pageRequest.pageSize

        var preIndex = -1L
        var curIndex: Long
        var total = 0L
        val result = ArrayList<E>()

        // 遍历所有分表进行查询
        val template = determineReactiveMongoOperations()
        for (sequence in 0 until shardingCount) {
            // 重置需要跳过的记录数量和limit
            query.skip(0L)
            query.limit(0)

            val collectionName = parseSequenceToCollectionName(sequence)

            // 统计总数
            val count = template.count(query, classType, collectionName).awaitSingle()
            if (count == 0L) {
                continue
            }
            total += count
            curIndex = total - 1

            // 当到达目标分页时才进行查询
            if (curIndex >= startIndex && limit > 0) {
                if (preIndex < startIndex) {
                    // 跳过属于前一个分页的数据
                    query.skip(startIndex - preIndex - 1)
                }
                query.limit(limit)
                val nodes = template.find(query, classType, collectionName).collectList().awaitSingle()
                // 更新还需要的数据数
                limit -= nodes.size
                result.addAll(nodes)
            }
            preIndex = curIndex
        }

        return PageImpl(result, pageRequest, total)
    }

    private fun shardingKeyToCollectionName(shardValues: List<Any>): String {
        return parseSequenceToCollectionName(shardingUtils.shardingSequenceFor(shardValues, shardingCount))
    }

    fun parseSequenceToCollectionName(sequence: Int): String {
        return collectionName + "_" + sequence
    }

    private fun filterExistedIndex(indexDefinitions: List<IndexDefinition>): List<IndexDefinition> {
        val reactiveMongoOperations = determineReactiveMongoOperations()
        val collectionName = parseSequenceToCollectionName(0)
        val indexInfoList = reactiveMongoOperations.indexOps(collectionName).indexInfo
        return indexInfoList.map { index -> index.name }.collectList().map { indexNameList ->
            val filteredList = indexDefinitions.filter { index ->
                val indexOptions = index.indexOptions
                if (indexOptions.contains("name")) {
                    val indexName = indexOptions.getString("name")
                    !indexNameList.contains(indexName)
                } else true
            }
            filteredList
        }.block() ?: emptyList()
    }

    abstract fun determineShardingUtils(): ShardingUtils

    companion object {
        private val logger = LoggerFactory.getLogger(ShardingMongoReactiveDao::class.java)
        private const val MAX_SHARDING_COUNT_OF_PAGE_QUERY = 256
    }
}
